package com.bruce.ai.alibaba.agentcard.discovery.auth;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Base64;
import java.util.Optional;

/**
 * JWT令牌验证器
 * 提供JWT令牌的解析、验证和信息提取功能
 * 
 * 特性：
 * - JWT格式验证
 * - 过期时间检查
 * - 声明提取
 * - 安全的Base64解码
 */
@Component
public class JwtTokenValidator {
    
    private static final Logger logger = LoggerFactory.getLogger(JwtTokenValidator.class);
    private final ObjectMapper objectMapper;
    
    public JwtTokenValidator() {
        this.objectMapper = new ObjectMapper();
    }
    
    /**
     * 验证JWT令牌
     * @param token JWT令牌
     * @return 验证结果
     */
    public TokenValidationResult validateToken(String token) {
        try {
            if (token == null || token.trim().isEmpty()) {
                return TokenValidationResult.invalid("Token is null or empty");
            }
            
            // 移除Bearer前缀（如果存在）
            if (token.startsWith("Bearer ")) {
                token = token.substring(7);
            }
            
            // 验证JWT格式
            String[] parts = token.split("\\.");
            if (parts.length != 3) {
                return TokenValidationResult.invalid("Invalid JWT format");
            }
            
            // 解析header
            JsonNode header = parseJwtPart(parts[0]);
            if (header == null) {
                return TokenValidationResult.invalid("Invalid JWT header");
            }
            
            // 解析payload
            JsonNode payload = parseJwtPart(parts[1]);
            if (payload == null) {
                return TokenValidationResult.invalid("Invalid JWT payload");
            }
            
            // 检查过期时间
            if (payload.has("exp")) {
                long exp = payload.get("exp").asLong();
                if (Instant.now().getEpochSecond() > exp) {
                    return TokenValidationResult.expired("Token has expired");
                }
            }
            
            // 检查生效时间
            if (payload.has("nbf")) {
                long nbf = payload.get("nbf").asLong();
                if (Instant.now().getEpochSecond() < nbf) {
                    return TokenValidationResult.invalid("Token not yet valid");
                }
            }
            
            return TokenValidationResult.valid(payload);
            
        } catch (Exception e) {
            logger.warn("Error validating JWT token", e);
            return TokenValidationResult.invalid("Token validation error: " + e.getMessage());
        }
    }
    
    /**
     * 从JWT令牌中提取主题（subject）
     * @param token JWT令牌
     * @return 主题，如果不存在则返回空
     */
    public Optional<String> extractSubject(String token) {
        TokenValidationResult result = validateToken(token);
        if (result.isValid() && result.getPayload().has("sub")) {
            return Optional.of(result.getPayload().get("sub").asText());
        }
        return Optional.empty();
    }
    
    /**
     * 从JWT令牌中提取发行者（issuer）
     * @param token JWT令牌
     * @return 发行者，如果不存在则返回空
     */
    public Optional<String> extractIssuer(String token) {
        TokenValidationResult result = validateToken(token);
        if (result.isValid() && result.getPayload().has("iss")) {
            return Optional.of(result.getPayload().get("iss").asText());
        }
        return Optional.empty();
    }
    
    /**
     * 从JWT令牌中提取受众（audience）
     * @param token JWT令牌
     * @return 受众，如果不存在则返回空
     */
    public Optional<String> extractAudience(String token) {
        TokenValidationResult result = validateToken(token);
        if (result.isValid() && result.getPayload().has("aud")) {
            return Optional.of(result.getPayload().get("aud").asText());
        }
        return Optional.empty();
    }
    
    /**
     * 检查令牌是否即将过期
     * @param token JWT令牌
     * @param thresholdSeconds 阈值秒数
     * @return 是否即将过期
     */
    public boolean isTokenExpiringSoon(String token, long thresholdSeconds) {
        TokenValidationResult result = validateToken(token);
        if (result.isValid() && result.getPayload().has("exp")) {
            long exp = result.getPayload().get("exp").asLong();
            long now = Instant.now().getEpochSecond();
            return (exp - now) <= thresholdSeconds;
        }
        return false;
    }
    
    /**
     * 解析JWT的某个部分
     * @param part Base64编码的JWT部分
     * @return 解析后的JSON节点
     */
    private JsonNode parseJwtPart(String part) {
        try {
            byte[] decodedBytes = Base64.getUrlDecoder().decode(part);
            String decodedString = new String(decodedBytes, StandardCharsets.UTF_8);
            return objectMapper.readTree(decodedString);
        } catch (Exception e) {
            logger.debug("Error parsing JWT part: {}", e.getMessage());
            return null;
        }
    }
    
    /**
     * JWT令牌验证结果
     */
    public static class TokenValidationResult {
        private final boolean valid;
        private final boolean expired;
        private final String errorMessage;
        private final JsonNode payload;
        
        private TokenValidationResult(boolean valid, boolean expired, String errorMessage, JsonNode payload) {
            this.valid = valid;
            this.expired = expired;
            this.errorMessage = errorMessage;
            this.payload = payload;
        }
        
        public static TokenValidationResult valid(JsonNode payload) {
            return new TokenValidationResult(true, false, null, payload);
        }
        
        public static TokenValidationResult invalid(String errorMessage) {
            return new TokenValidationResult(false, false, errorMessage, null);
        }
        
        public static TokenValidationResult expired(String errorMessage) {
            return new TokenValidationResult(false, true, errorMessage, null);
        }
        
        public boolean isValid() {
            return valid;
        }
        
        public boolean isExpired() {
            return expired;
        }
        
        public String getErrorMessage() {
            return errorMessage;
        }
        
        public JsonNode getPayload() {
            return payload;
        }
    }
}